import json
import os
import sys

sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from aws_model.aws_glue import Operation_aws_glue
from aws_model.athena_submit import submit
from aws_model.aws_stepfunc import send_task_failure, send_task_success
from base_class.inspect_logging import print_to_log
from aws_model.aws_s3 import Operation_aws_s3

s3 = Operation_aws_s3()
glue = Operation_aws_glue()


def run(table_name: str, db: str):
    table_desc = glue.glue.get_table(
        DatabaseName=db,
        Name=table_name,
    )
    location = table_desc['Table']['StorageDescriptor']['Location']
    partition_keys = []
    for partition in table_desc['Table']['PartitionKeys']:
        partition_keys.append(partition['Name'])
    files = s3.ls(location)
    file_partitions = {}
    for file in files:
        paths = file['object'].split('/')
        partition_names = []
        for partition_key in partition_keys:
            for one in paths:
                if partition_key in one:
                    partition_names.append(one)
        d = file_partitions
        for partition_name in partition_names:
            name = partition_name.split('=')[1]
            if not name in d:
                d[name] = {}
            d = d[name]

    exists_partitions = glue.get_partitions(db=db, table_name=table_name)

    drop_partitions = []
    for one in exists_partitions:
        d = file_partitions
        fit = True
        for p_k in one:
            if not p_k in d:
                drop_partitions.append(one)
                fit = False
                break
            else:
                d = d[p_k]
        if fit:
            d = file_partitions
            for p_k in one:
                if len(d[p_k]) == 0:
                    del d[p_k]
                else:
                    d = d[p_k]
    add_partitions = []
    for key in file_partitions:
        dict_to_list(file_partitions[key], key, add_partitions, len(partition_keys))
    modify_partitions(db, table_name, drop_partitions, partition_keys, 'DROP')
    modify_partitions(db, table_name, add_partitions, partition_keys, 'ADD')
    output = {"msg": f"{table_name} submit ok", "result": "success"}
    output = json.dumps(output)
    resp = send_task_success(tasktoken, output)
    print_to_log("step_response: ", resp)


def dict_to_list(d, s, li, size):
    if len(d) == 0:
        if len(s.split(',')) == size:
            li.append(s.split(','))
    else:
        for key in d:
            dict_to_list(d[key], s + "," + key, li, size)


def modify_partitions(db, table_name, partitions, partition_keys, add_or_drop):
    sql = f'ALTER TABLE {db}.{table_name} {add_or_drop} '
    i = 0
    for partition in partitions:
        i += 1
        filters = []
        for j in range(0, len(partition_keys)):
            filters.append(f'''{partition_keys[j]} = '{partition[j]}' ''')
        filter_sql = ','.join(filters)
        sql = f'{sql} PARTITION ({filter_sql}) '
        if add_or_drop.lower() == 'drop':
            sql = sql + ','
        if i % 20 == 0:
            print(add_or_drop, i, '个分区')
            if add_or_drop.lower() == 'drop':
                sql = sql[0:sql.rfind(',')]
            print(sql)
            submit(sql)
            sql = f'ALTER TABLE {db}.{table_name} {add_or_drop} '

    if 'PARTITION' in sql:
        print(add_or_drop, i, '个分区')
        if add_or_drop.lower() == 'drop':
            sql = sql[0:sql.rfind(',')]
        print(sql)
        submit(sql)


if __name__ == '__main__':
    try:
        event = json.loads(sys.argv[1])
        tasktoken = event.get('TaskToken')
        # table_name = 'ac_mqtt_pass_through_hour'
        # db = 'original'
        table_name = event.get('table_name')
        db = event.get('db')
        run(table_name=table_name, db=db)
    except Exception as e:
        print_to_log("任务 失败")
        print_to_log("原因： ", str(e))
        send_task_failure(str(e), f'执行{sys._getframe().f_code.co_filename}错误', tasktoken)
